#!/usr/bin/env bash

declare -g stress_tool="stress-ng"
declare -g network_tool="iperf3"
declare -g stress_file="/tmp/collect/stress_info"
declare -g stress_usage="/tmp/collect/stress_usage"

function stress_checker() {
    which $stress_tool 2>&1 > /dev/null
    if [ $? -ne 0 ]; then
        return 1
    fi

    touch ${stress_file}
    touch ${stress_usage}

    echo -e "begin-timestamp\tend-timestamp\ttype\tstress\tcommand" >> ${stress_file}
    return 0
}

function execute_stress_command() {
    local type=$1
    local stress=$2
    local command=$3

    btimestamp=$(date "+%Y-%m-%d %H:%M:%S")
    eval $command
    etimestamp=$(date "+%Y-%m-%d %H:%M:%S")

    echo -e "${btimestamp}\t${etimestamp}\t${type}\t${stress}\t${command}" >> $stress_file

    sleep 5s
}

function collect_stress_usage() {
    local timestamp=$(date "+%Y-%m-%d %H:%M:%S")
    # 100 means 1 core
    local cpu_percent=$(ps -aux | grep -E 'stress-ng | iperf3' | grep -v grep | awk '{sum +=$3}END{print sum}')
    if [ "x$cpu_percent" == x ]; then
        cpu_percent=0
    fi

    local memory_percent=$(ps -aux | grep 'stress-ng | iperf3' | grep -v grep | awk '{sum +=$4}END{print sum}')
    if [ "x$memory_percent" == x ]; then
        memory_percent=0
    fi

    pid=$(ps -aux | grep 'stress-ng | iperf3' | grep -v grep | awk '{print $2}')
    for p in $pid; do
        args="-p $p"
    done

    local disk_read=$(iotop -b -o -n 1 $args | grep "Total DISK READ" | tr ":" " " | awk '{print $4 $5}')
    local disk_write=$(iotop -b -o -n 1 $args | grep "Total DISK WRITE" | tr ":" " " |awk '{print $10 $11}')

    echo -e "${timestamp}\t${cpu_percent}\t${memory_percent}\t${disk_read}\t${disk_write}" >> ${stress_usage}
}

function collect_stress_usage_loop() {
    while true; do
        collect_stress_usage
        sleep 1
    done
}

function no_stress() {
    local type=$1

    local command="sleep 30s"
    execute_stress_command "$type" "0" "$command"
}

function stress_cpu() {
    local total_cpu=$(cat /proc/cpuinfo | grep "processor" | wc -l)
    local max_cpu=$(round $(echo "${total_cpu}*0.7" | bc) 0)

    no_stress "cpu"

    for cpu_num in $(seq 1 $max_cpu); do
        local command="${stress_tool} --quiet --timeout 30s --cpu $cpu_num"

        execute_stress_command "cpu" "$cpu_num" "$command"
    done
}

function stress_memory() {
    local available_memory=$(cat /proc/meminfo | grep MemAvailable | awk '{print $2}')
    local max_mem_percent=0.7

    no_stress "memory"

    for mem_percent in $(seq 0.05 0.05 $max_mem_percent); do
        local memory_bytes=$(echo ${available_memory}*${mem_percent} | bc)

        local command="${stress_tool} --quiet --timeout 30s --vm 1 --vm-bytes ${memory_bytes}k"
        execute_stress_command "memory" "${memory_bytes}" "$command"
    done
}

function stress_cache() {
    local cache_level=(1 2 3)
    local total_cpu=$(cat /proc/cpuinfo | grep "processor" | wc -l)

    for level in ${cache_level[*]}; do

        no_stress "L${level} cache"

        for cpu_num in $(seq 1 $total_cpu); do
            local command="${stress_tool} --quiet --timeout 30s --cache $cpu_num --cache-level ${level}"
            execute_stress_command "L${level} cache" "$cpu_num" "$command"
        done
    done

    uname -a | grep -E 'x86|amd'
    if [ $? -eq 0 ]; then

        no_stress "icache"

        for cpu_num in $(seq 1 $total_cpu); do
            local command="${stress_tool} --quiet --timeout 30s --icache $cpu_num"
            execute_stress_command "icache" "$cpu_num" "$command"
        done
    fi
}

function stress_memory_bandwidth() {
    local max_bandwidth=8000

    no_stress "mem-bandwidth"

    for mem_bandwidth in $(seq 500 500 $max_bandwidth); do
        local command="${stress_tool} --quiet --timeout 30s --memrate 1 --memrate-bytes ${mem_bandwidth}M"
        execute_stress_command "mem-bandwidth" "${mem_bandwidth}" "$command"
    done
}

function stress_disk_io() {
    # stress-ng will push I/O as fas as it can.
    local disk_location=$1
    local max_disk_io=1000
    local current_dir=$(cd $(dirname $0) && pwd)

    cd $disk_location

    no_stress "disk io"

    for disk_io in $(seq 50 50 $max_disk_io); do
        local command="${stress_tool} --quiet --timeout 30s --io 1 --hdd 1 --hdd-bytes ${disk_io}M"
        execute_stress_command "disk io" "$disk_io" "$command"
    done

    cd $current_dir
}

function stress_network() {
    local ip=$1
    local port=$2

    if [ "x$ip" == "x0" ]; then
        return
    fi

    if [ "x$port" == "x0" ]; then
        port=5201
    fi

    which $network_tool 2>&1 > /dev/null
    if [ $? -ne 0 ]; then
        echo "$network_tool not found"
        return
    fi

    # server command: iperf3 -s
    local max_network_bandwidth=1000

    no_stress "network out"

    for bandwidth in $(seq 50 50 $max_network_bandwidth); do
        local command="$network_tool -t 30s -u -c $ip -p $port -i 1 -b ${bandwidth}M"
        execute_stress_command "network out" "${bandwidth}" "$command"
    done

    no_stress "network in"

    for bandwidth in $(seq 50 50 $max_network_bandwidth); do
        local command="$network_tool -t 30s -u -c $ip -p $port -i 1 -R -b ${bandwidth}M"
        execute_stress_command "network in" "${bandwidth}" "$command"
    done
}

function stress_cpu_cache_mixed() {
    local total_cpu=$(cat /proc/cpuinfo | grep "processor" | wc -l)
    local max_cpu=$(round $(echo "${total_cpu}*0.7" | bc) 0)

    no_stress "mixed-cpu-cache"

    for cache_num in $(seq 1 $total_cpu); do
        local command="${stress_tool} --quiet --timeout 30s --cache $cache_num --cache-level 1"
        command="${command} & ${stress_tool} --quiet --timeout 30s --cache $cache_num --cache-level 2"
        command="${command} & ${stress_tool} --quiet --timeout 30s --cache $cache_num --cache-level 3"

        uname -a | grep -E 'x86|amd'
        if [ $? -eq 0 ]; then
            command="${command} & ${stress_tool} --quiet --timeout 30s --icache $cache_num"
        fi

        cpu_num=$cache_num
        if [ $cpu_num -gt $max_cpu ]; then
            cpu_num=$max_cpu
        fi
        
        command="${command} & ${stress_tool} --quiet --timeout 30s --cpu $cpu_num"

        execute_stress_command "mixed-cpu-cache" "$cache_num" "$command"
    done
}

function round() {
    printf "%.${2}f" "${1}"
}
